package com.seewo.dubbo.debugger.processor;


import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.seewo.dubbo.debugger.common.AnnotatedClass;
import com.seewo.dubbo.debugger.common.AnnotatedMethod;
import com.seewo.dubbo.debugger.common.annotation.Api;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.TypeElement;
import javax.lang.model.util.Elements;
import javax.tools.Diagnostic;
import javax.tools.FileObject;
import javax.tools.StandardLocation;
import java.io.IOException;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * @author llchen12
 * @date 2018/6/26
 */
@SupportedAnnotationTypes({"com.seewo.dubbo.debugger.common.annotation.Api"})
public class DocCommentProcessor extends AbstractProcessor {

    private static final String METADATA_PATH = "META-INF/doc-comment.json";


    private Elements elementUtils;
    private Messager messager;
    private Filer filer;

    private Map<String, AnnotatedClass> docClassMap = new HashMap<>(64);

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        this.elementUtils = processingEnv.getElementUtils();
        this.messager = processingEnv.getMessager();
        this.filer = processingEnv.getFiler();
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {

        Set<? extends Element> apiEleSet = roundEnv.getElementsAnnotatedWith(Api.class);

        processDocMethods(apiEleSet);

        if (roundEnv.processingOver()) {
            try {
                writeData();
            } catch (IOException e) {
                messager.printMessage(Diagnostic.Kind.WARNING, "Failed to write metadata " + e.getMessage());
            }
        }
        return false;
    }

    private void processDocMethods(Set<? extends Element> docMethodSet) {
        for (Element ele : docMethodSet) {
            if (ele.getKind() != ElementKind.METHOD) {
                continue;
            }
            AnnotatedMethod method = new AnnotatedMethod((ExecutableElement) ele, elementUtils.getDocComment(ele));
            String className = method.getClassName();
            AnnotatedClass aClass = docClassMap.get(className);
            if (aClass == null) {
                Element parent = ele.getEnclosingElement();
                aClass = new AnnotatedClass(className, elementUtils.getDocComment(parent));
                aClass.addMethod(method);
                docClassMap.put(className, aClass);
            } else {
                aClass.addMethod(method);
            }
        }
    }

    private void processDocClasses(Set<Element> docClassSet) {
        for (Element ele : docClassSet) {
            if (ele.getKind() != ElementKind.CLASS) {
                continue;
            }
            String className = ((TypeElement) ele).getQualifiedName().toString();
            if (!docClassMap.containsKey(className)) {
                AnnotatedClass annotatedClass = new AnnotatedClass(className, elementUtils.getDocComment(ele));
                docClassMap.put(className, annotatedClass);
            }
        }
    }


    private void writeData() throws IOException {
        if (docClassMap.isEmpty()) {
            return;
        }
        FileObject resource = filer.createResource(StandardLocation.CLASS_OUTPUT, "", METADATA_PATH);
        try (OutputStream outputStream = resource.openOutputStream()) {
            String json = JSON.toJSONString(docClassMap, SerializerFeature.PrettyFormat,
                    SerializerFeature.SkipTransientField);
            outputStream.write(json.getBytes("UTF-8"));
        }
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }
}
